#pragma once
#include "acl/acl.h"
#include "../include/custom_type.h"
#include "./data_utils.h"
#include <cstdio>
#include <string>
#include <algorithm>

/**
 * @brief: main函数中调用host的工具函数
 * @param [in] kernel_host: 本次调用的host函数
 * @param [in] layoutA: A矩阵排布格式
 * @param [in] layoutB: B矩阵排布格式
 * @param [in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param [in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param [in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param [in] batchCount: 批量矩阵乘的batch数
 * @param [in] path_A: 本次GEMM计算的A矩阵数据（如果有）读取路径
 * @param [in] path_B: 本次GEMM计算的B矩阵数据（如果有）读取路径
 * @param [in] path_C: 本次GEMM计算的C矩阵数据（如果有）读取路径
 * @param [in] path_alpha: 本次GEMM计算的alpha数据（如果有）读取路径
 * @param [in] path_beta:  本次GEMM计算的beta 数据（如果有）读取路径
 * @param [in] path_validM: 本次GEMM计算的validM数据（如果有）读取路径（host中以zeroPaddingMNK validM来生成4个任务的validNK）
 * @param [in] pre_output: 上一任务的输出（如果有），作为本次任务的输入，最后会free掉
 * @param [out] cur_output: 本次任务的输出，作为下一任务的输入（如果有），最后会返回
 * @param [in] mode: 0-GroupGemm/strideBatchedGemm的满zeroPaddingMNK矩阵计算，123-三个LLMsGEMM任务的有效矩阵维度计算
 * @param [in] test_mode: 是否需要读取和验证真值数据。0-功能测试，1-性能测试
 */
void run_main(
    void (* kernel_host)(
        layoutType,         // layoutA
        layoutType,         // layoutB
        int64_t,            // zeroPaddingM
        int64_t,            // zeroPaddingN
        int64_t,            // zeroPaddingK
        int64_t,            // batchCount
        int64_t *,          // d_validM
        half,               // alpha
        half *,             // d_A
        half *,             // d_B
        half,               // beta
        half *,             // d_C
        aclrtStream &,      // stream
        uint8_t             // mode
    ), 
    aclrtStream &stream, 
    layoutType layoutA, 
    layoutType layoutB, 
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN, 
    int64_t zeroPaddingK, 
    int64_t batchCount, 
    std::string &path_A, 
    std::string &path_B, 
    std::string &path_C, 
    std::string &path_alpha, 
    std::string &path_beta, 
    std::string &path_validM, 
    std::string &path_expect_result, 
    half* &pre_output,          
    half* &cur_output, 
    uint8_t mode, 
    uint8_t test_mode
){

    // test_mode == 0: 需要分配空间后，读取脚本生成的golden输入数据，然后将计算结果与脚本生成的golden输出数据进行对比，验证功能
    // test_mode == 1: 不读取脚本生成数据验证结果，仅分配空间，然后计算，用于测试性能。不随机有效矩阵数据大小，都用满的zeroPaddingM*N*K数据来计算
    // uint8_t test_mode = 0;

    if(test_mode == 0){ // 验证功能
        // 计算输入矩阵的字节数用于分配空间
        uint64_t inputABytes = zeroPaddingM * zeroPaddingK * sizeof( half ) * batchCount;
        uint64_t inputBBytes = zeroPaddingK * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputCBytes = zeroPaddingM * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputAlphaBetaBytes = sizeof( half );
        uint64_t inputValidMBytes = sizeof( int64_t ) * batchCount;
        uint64_t inputValidNBytes = sizeof( int64_t ) * batchCount;
        uint64_t inputValidKBytes = sizeof( int64_t ) * batchCount;

        // 定义host侧指针
        half *h_A = nullptr;
        half *h_B = nullptr;
        half *h_C = nullptr;
        half *h_alpha = nullptr;
        half *h_beta = nullptr;
        int64_t *h_validM = nullptr;

        // aclrtMallocHost 分配host侧内存
        if(pre_output == nullptr)
            ACL_CHECK(aclrtMallocHost((void **)&h_A, inputABytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_B, inputBBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_C, inputCBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_alpha, inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_beta , inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_validM, inputValidMBytes));
        // ACL_CHECK(aclrtMallocHost((void **)&h_validN, inputValidNBytes));
        // ACL_CHECK(aclrtMallocHost((void **)&h_validK, inputValidKBytes));

        // 读入Python脚本生成的测试输入数据到host侧内存
        if(pre_output == nullptr)
            ReadFile( path_A, h_A, inputABytes );
        ReadFile( path_B, h_B, inputBBytes );
        ReadFile( path_C, h_C, inputCBytes );
        ReadFile( path_alpha, h_alpha, inputAlphaBetaBytes );
        ReadFile( path_beta , h_beta , inputAlphaBetaBytes );
        ReadFile( path_validM, h_validM, inputValidMBytes );


        // test input
        // printf("test input data: \n");
        // printf("\n");
        // printf("validM: \n");
        // for(int64_t i = 0; i < batchCount; i++){
        //     printf("%ld ", h_validM[i]);
        // }
        // printf("\n");
        // printf("\n");

        // printf("A:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingK; j++){
        //         printf("%f ", h_A[i*zeroPaddingK+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // printf("B:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingK; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_B[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // printf("C:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_C[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        printf("\n");
        printf("alpha = %f, beta = %f\n", *h_alpha, *h_beta);
        printf("\n");
        // test input finished

        // 定义device侧指针
        half *d_A = nullptr;
        half *d_B = nullptr;
        half *d_C = nullptr;
        int64_t *d_validM = nullptr;
        half alpha = 0.0;
        half beta  = 0.0;

        // aclrtMalloc 分配device侧内存空间
        if(pre_output == nullptr)
            ACL_CHECK(aclrtMalloc((void **)&d_A, inputABytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_B, inputBBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_C, inputCBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_validM, inputValidMBytes, ACL_MEM_MALLOC_HUGE_FIRST));

        // alpha 和 beta 直接取值为常量
        alpha = *h_alpha;
        beta  = *h_beta;
        // alpha = 1.0f;
        // beta = 0.0f;
        
        // aclrtMemcpy 把host侧数据copy到device侧
        if(pre_output == nullptr)
            ACL_CHECK(aclrtMemcpy(d_A, inputABytes, h_A, inputABytes, ACL_MEMCPY_HOST_TO_DEVICE));
        else
            d_A = pre_output;
        ACL_CHECK(aclrtMemcpy(d_B, inputBBytes, h_B, inputBBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_C, inputCBytes, h_C, inputCBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_validM, inputValidMBytes, h_validM, inputValidMBytes, ACL_MEMCPY_HOST_TO_DEVICE));

        // 算子host侧调用
        (*kernel_host)(
            layoutA, 
            layoutB, 
            zeroPaddingM, 
            zeroPaddingN, 
            zeroPaddingK, 
            batchCount, 
            d_validM, 
            alpha, 
            d_A, 
            d_B, 
            beta, 
            d_C, 
            stream, 
            mode
        );

        // 计算结果读回host侧，与期望数据做比较
        ACL_CHECK(aclrtMemcpy(h_C, inputCBytes, d_C, inputCBytes, ACL_MEMCPY_DEVICE_TO_HOST));

        // 读取期望数据
        half *h_expect_result = nullptr;
        ACL_CHECK(aclrtMallocHost((void **)&h_expect_result, inputCBytes));
        ReadFile( path_expect_result, h_expect_result, inputCBytes );

        // test output
        // printf("test output data: \n");
        // printf("\n");

        // printf("expect_result:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_expect_result[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");

        // printf("output_C:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_C[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // test output finished

        // 对比理想结果与计算结果，计算精度、误差等
        calOutputError<half>( h_C, h_expect_result, zeroPaddingM * zeroPaddingN * batchCount, zeroPaddingK, zeroPaddingM, zeroPaddingN );

        // 释放host侧内存
        if(pre_output == nullptr)
            ACL_CHECK(aclrtFreeHost(h_A));
        ACL_CHECK(aclrtFreeHost(h_B));
        ACL_CHECK(aclrtFreeHost(h_C));
        ACL_CHECK(aclrtFreeHost(h_alpha));
        ACL_CHECK(aclrtFreeHost(h_beta));
        ACL_CHECK(aclrtFreeHost(h_validM));
        ACL_CHECK(aclrtFreeHost(h_expect_result));

        // 释放device侧内存
        ACL_CHECK(aclrtFree(d_A));
        ACL_CHECK(aclrtFree(d_B));
        // ACL_CHECK(aclrtFree(d_C));
        cur_output = d_C;
        ACL_CHECK(aclrtFree(d_validM));

    }else if(test_mode == 1){// 不验证功能 用满的zeroPadingM*N*K的数据来计算

        // 计算输入矩阵的字节数用于分配空间
        uint64_t inputABytes = zeroPaddingM * zeroPaddingK * sizeof( half ) * batchCount;
        uint64_t inputBBytes = zeroPaddingK * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputCBytes = zeroPaddingM * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputAlphaBetaBytes = sizeof( half );
        uint64_t inputValidMBytes = sizeof( int64_t ) * batchCount;
        uint64_t inputValidNBytes = sizeof( int64_t ) * batchCount;
        uint64_t inputValidKBytes = sizeof( int64_t ) * batchCount;

        // 定义host侧指针
        half *h_A = nullptr;
        half *h_B = nullptr;
        half *h_C = nullptr;
        half *h_alpha = nullptr;
        half *h_beta = nullptr;
        int64_t *h_validM = nullptr;

        // aclrtMallocHost 分配host侧内存
        if(pre_output == nullptr)
            ACL_CHECK(aclrtMallocHost((void **)&h_A, inputABytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_B, inputBBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_C, inputCBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_alpha, inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_beta , inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_validM, inputValidMBytes));

        // 不再读取Python脚本生成的golden数据
        // if(pre_output == nullptr)
        //     ReadFile( path_A, h_A, inputABytes );
        // ReadFile( path_B, h_B, inputBBytes );
        // ReadFile( path_C, h_C, inputCBytes );
        // ReadFile( path_alpha, h_alpha, inputAlphaBetaBytes );
        // ReadFile( path_beta , h_beta , inputAlphaBetaBytes );
        // ReadFile( path_validM, h_validM, inputValidMBytes );
  
        std::fill(h_validM, h_validM + batchCount, zeroPaddingM);

        // test input
        // printf("test input data: \n");
        // printf("\n");
        // printf("validM: \n");
        // for(int64_t i = 0; i < batchCount; i++){
        //     printf("%ld ", h_validM[i]);
        // }
        // printf("\n");
        // printf("\n");

        // printf("A:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingK; j++){
        //         printf("%f ", h_A[i*zeroPaddingK+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // printf("B:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingK; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_B[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // printf("C:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_C[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // printf("alpha = %f, beta = %f\n", *h_alpha, *h_beta);
        // printf("\n");
        // test input finished

        // 定义device侧指针
        half *d_A = nullptr;
        half *d_B = nullptr;
        half *d_C = nullptr;
        int64_t *d_validM = nullptr;
        half alpha = 0.0;
        half beta  = 0.0;

        // aclrtMalloc 分配device侧内存空间
        if(pre_output == nullptr)
            ACL_CHECK(aclrtMalloc((void **)&d_A, inputABytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_B, inputBBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_C, inputCBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_validM, inputValidMBytes, ACL_MEM_MALLOC_HUGE_FIRST));

        // alpha 和 beta 直接取定值
        alpha = 0.5f;
        beta  = 0.5f;
        // alpha = 1.0f;
        // beta = 0.0f;
        
        printf("\n");
        printf("alpha = %f, beta = %f\n", alpha, beta);
        printf("\n");

        // aclrtMemcpy 把host侧数据copy到device侧
        if(pre_output == nullptr)
            ACL_CHECK(aclrtMemcpy(d_A, inputABytes, h_A, inputABytes, ACL_MEMCPY_HOST_TO_DEVICE));
        else
            d_A = pre_output;
        ACL_CHECK(aclrtMemcpy(d_B, inputBBytes, h_B, inputBBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_C, inputCBytes, h_C, inputCBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_validM, inputValidMBytes, h_validM, inputValidMBytes, ACL_MEMCPY_HOST_TO_DEVICE));

        // 算子host侧调用
        (*kernel_host)(
            layoutA, 
            layoutB, 
            zeroPaddingM, 
            zeroPaddingN, 
            zeroPaddingK, 
            batchCount, 
            d_validM, 
            alpha, 
            d_A, 
            d_B, 
            beta, 
            d_C, 
            stream, 
            0 // 只能是满的矩阵
        );

        // 不需要比较数据
        // ACL_CHECK(aclrtMemcpy(h_C, inputCBytes, d_C, inputCBytes, ACL_MEMCPY_DEVICE_TO_HOST));

        // 读取期望数据
        // half *h_expect_result = nullptr;
        // ACL_CHECK(aclrtMallocHost((void **)&h_expect_result, inputCBytes));
        // ReadFile( "./data/expect_result.bin", h_expect_result, inputCBytes );

        // test output
        // printf("test output data: \n");
        // printf("\n");

        // printf("expect_result:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_expect_result[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");

        // printf("output_C:\n");
        // for(int64_t i = 0; i < batchCount * zeroPaddingM; i++ ){
        //     for(int64_t j = 0; j < zeroPaddingN; j++){
        //         printf("%f ", h_C[i*zeroPaddingN+j]);
        //     }
        //     printf("\n");
        // }
        // printf("\n");
        // test output finished

        // 不需要对比理想结果与计算结果
        // calOutputError<half>( h_C, h_expect_result, inputCBytes / sizeof(half), zeroPaddingK, zeroPaddingM, zeroPaddingN );

        // 释放host侧内存
        if(pre_output == nullptr)
            ACL_CHECK(aclrtFreeHost(h_A));
        ACL_CHECK(aclrtFreeHost(h_B));
        ACL_CHECK(aclrtFreeHost(h_C));
        ACL_CHECK(aclrtFreeHost(h_alpha));
        ACL_CHECK(aclrtFreeHost(h_beta));
        ACL_CHECK(aclrtFreeHost(h_validM));
        // ACL_CHECK(aclrtFreeHost(h_expect_result));

        // 释放device侧内存
        ACL_CHECK(aclrtFree(d_A));
        ACL_CHECK(aclrtFree(d_B));
        // ACL_CHECK(aclrtFree(d_C));
        cur_output = d_C;
        ACL_CHECK(aclrtFree(d_validM));

    }


    



}